package net.sudot.commons.security;

import net.sudot.commons.utils.WebUtils;
import org.apache.shiro.authz.AuthorizationException;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;

/**
 * CSRF拦截器
 *
 * @author tangjialin on 2018-03-17.
 */
public class CsrfInterceptor extends HandlerInterceptorAdapter {
    /**
     * "CSRF令牌"Parameter,Attribute"名称
     */
    private static final String CSRF_TOKEN_NAME = "CSRF_TOKEN";
    /**
     * "CSRF令牌"Header"名称
     */
    private static final String CSRF_TOKEN_HEADER_NAME = "X-XSRF-TOKEN";
    /**
     * "CSRF令牌"Cookie"名称
     */
    private static final String CSRF_TOKEN_COOKIE_NAME = "XSRF-TOKEN";

    /**
     * 默认无需防护的请求方法
     */
    private static final Set<String> DEFAULT_NOT_REQUIRE_PROTECTION_REQUEST_METHODS = new HashSet<String>();

    static {
        DEFAULT_NOT_REQUIRE_PROTECTION_REQUEST_METHODS.add("GET");
    }

    /**
     * 无需防护的请求方法
     */
    private Set<String> notRequireProtectionRequestMethods = DEFAULT_NOT_REQUIRE_PROTECTION_REQUEST_METHODS;

    /**
     * 请求前处理
     *
     * @param request  HttpServletRequest
     * @param response HttpServletResponse
     * @param handler  处理器
     * @return 是否继续执行
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Ignore ignore = handlerMethod.getMethodAnnotation(Ignore.class);
            if (ignore != null) { return true; }
            String csrfToken = WebUtils.getCookie(request, CSRF_TOKEN_COOKIE_NAME);
            if (csrfToken == null || csrfToken.isEmpty()) {
                csrfToken = UUID.randomUUID().toString();
                WebUtils.addCookie(request, response, CSRF_TOKEN_COOKIE_NAME, csrfToken);
            }
            String method = request.getMethod();
            if (method != null && !notRequireProtectionRequestMethods.contains(method.toUpperCase())) {
                String actualCsrfToken = request.getParameter(CSRF_TOKEN_NAME);
                if (actualCsrfToken == null) {
                    actualCsrfToken = request.getHeader(CSRF_TOKEN_HEADER_NAME);
                }
                if (!csrfToken.equals(actualCsrfToken)) {
                    throw new AuthorizationException("无效的CSRF令牌");
                }
            }
            request.setAttribute(CSRF_TOKEN_NAME, csrfToken);
        }
        return super.preHandle(request, response, handler);
    }

    /**
     * 获取无需防护的请求方法
     *
     * @return 无需防护的请求方法
     */
    public Set<String> getNotRequireProtectionRequestMethods() {
        return notRequireProtectionRequestMethods;
    }

    /**
     * 设置无需防护的请求方法
     *
     * @param notRequireProtectionRequestMethods 无需防护的请求方法
     */
    public CsrfInterceptor setNotRequireProtectionRequestMethods(Set<String> notRequireProtectionRequestMethods) {
        this.notRequireProtectionRequestMethods = notRequireProtectionRequestMethods;
        return this;
    }

    /**
     * 控制器上加此注解,表示可不受CSRF拦截器检测
     *
     * @author tangjialin on 2018-04-19.
     */
    @Target({ElementType.METHOD})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface Ignore {
    }

}